data{
  int<lower=0> N;
  int<lower=1> K; //number of category
  int<lower=1> M; // number of covariates
  int<lower=2> F; // number of focus category
  matrix[N, M] covariates;
	int<lower=0, upper=F-1> focus[N];
  int<lower=0, upper=1> treatment[N];
  int<lower=1,upper=K> Y[N];
}
parameters{
  vector[M] betaX; // beta for covariates
  real beta; // beta for treatment
	real beta_f; // beta for focus
	real beta_inter; // beta for interaction term
  ordered[K-1] c;
}

transformed parameters{
	real<lower=0> Count[F];
	real id;

	//initialize
	for(f in 1:F)
		Count[f] = 0;

	for(i in 1:N){
		id = focus[i] + 1;
		for(f in 1:F){
			if(id == f)
				Count[f] = Count[f] + 1;
		}
	}

}

model{
  for(i in 1:N){
    Y[i] ~ ordered_logistic(covariates[i,] * betaX + beta*treatment[i] + beta_f*focus[i] + beta_inter*treatment[i]*focus[i], c);
  }

}
generated quantities{
	matrix[K,F] T;
	matrix[K,F] C;
	matrix[K,F] diff;
	real temp1;
	real temp2;

  matrix[2,F] diff_D;  // dichotomized
  matrix[2,F] diff_DwN;  // dichotomoized with neutral

	for(k in 1:K){
		for(f in 1:F){
			T[k,f]=0;
			C[k,f]=0;
			diff[k,f]=0;
		}
	}

	for(k in 1:2){
		for(f in 1:F){
			diff_D[k,f]=0;
			diff_DwN[k,f]=0;
		}
	}

	for(i in 1:N){
		for(f in 1:F){
			if(focus[i] == f-1){

				T[1,f] = inv_logit(- covariates[i,] * betaX - beta*1 - beta_f * focus[i] - beta_inter*1*focus[i] + c[1]);
				C[1,f] = inv_logit(- covariates[i,] * betaX - beta*0 - beta_f * focus[i] - beta_inter*0*focus[i] + c[1]);	
				diff[1,f] = diff[1,f] + T[1,f] - C[1,f];

				// Note:
				// In Stan: math/stan/math/prim/scal/fun/inv_logit.hpp
				// The inverse logit function is defined by
				// \mbox{logit}^{-1}(x) = \frac{1}{1 + \exp(-x)}
				// We gonna input x, so take -
				// https://www3.nd.edu/~rwilliam/xsoc73994/Ologit01.pdf <- correct?
				// http://docs.zeligproject.org/articles/zeligchoice_ologit.html
				// 
				// Depending on the value of x, either simplified (one exp) or complex (two exps) version is used in Stan
				// Then, the issue is depending on which one to use, the order of taking difference changes
				// That's why we need to take max and min

				for(k in 2:K-1){

					temp1 = inv_logit(- covariates[i,] * betaX - beta*1 - beta_f * focus[i] - beta_inter*1*focus[i] + c[k-1]);
					temp2 = inv_logit(- covariates[i,] * betaX - beta*1 - beta_f * focus[i] - beta_inter*1*focus[i] + c[k]);
					T[k,f] = fmax(temp1, temp2) - fmin(temp1, temp2);


					temp1 = inv_logit(- covariates[i,] * betaX - beta*0 - beta_f * focus[i] - beta_inter*0*focus[i] + c[k-1]);
					temp2 = inv_logit(- covariates[i,] * betaX - beta*0 - beta_f * focus[i] - beta_inter*0*focus[i] + c[k]);
					C[k,f] = fmax(temp1, temp2) - fmin(temp1, temp2); 
					diff[k,f] = diff[k,f] + T[k,f] - C[k,f];
				}

				T[K,f] =  1 - inv_logit(- covariates[i,] * betaX - beta*1 - beta_f * focus[i] - beta_inter*1*focus[i] + c[K-1]);
				C[K,f] =  1 - inv_logit(- covariates[i,] * betaX - beta*0 - beta_f * focus[i] - beta_inter*0*focus[i] + c[K-1]);
				diff[K,f] = diff[K,f] + T[K,f] - C[K,f];


        diff_D[1,f] = diff[1,f] + diff[2,f];
        diff_D[2,f] = diff[4,f] + diff[5,f];

        diff_DwN[1,f] = diff[1,f] + diff[2,f];
        diff_DwN[2,f] = diff[3,f] + diff[4,f] + diff[5,f];

			}
		}

	}

	for(k in 1:K){
		for(f in 1:F){
			diff[k,f] = diff[k,f] / Count[f];
		}
	}


	for(k in 1:2){
		for(f in 1:F){
			diff_D[k,f] = diff_D[k,f] / Count[f];
			diff_DwN[k,f] = diff_DwN[k,f] / Count[f];
		}
	}

}


